use std::{borrow::Cow, io::Write, ops::Deref, sync::Arc};

use anyhow::Result;
use bincode::{
    Decode, Encode,
    de::Decoder,
    enc::Encoder,
    error::{DecodeError, EncodeError},
};
use bytes_str::BytesStr;
use either::Either;
use once_cell::sync::Lazy;
use ref_cast::RefCast;
use regex::Regex;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use swc_sourcemap::{DecodedMap, SourceMap as RegularMap, SourceMapBuilder, SourceMapIndex};
use turbo_rcstr::{RcStr, rcstr};
use turbo_tasks::{ResolvedVc, TryJoinIterExt, Vc};
use turbo_tasks_fs::{
    File, FileContent, FileSystem, FileSystemPath, VirtualFileSystem,
    rope::{Rope, RopeBuilder},
};

use crate::{
    SOURCE_URL_PROTOCOL, asset::AssetContent, source::Source,
    source_map::utils::add_default_ignore_list, source_pos::SourcePos,
    virtual_source::VirtualSource,
};

pub(crate) mod source_map_asset;
pub mod utils;

pub use source_map_asset::SourceMapAsset;

/// Represents an empty value in a u32 variable in the sourcemap crate.
static SOURCEMAP_CRATE_NONE_U32: u32 = !0;

/// Allows callers to generate source maps.
#[turbo_tasks::value_trait]
pub trait GenerateSourceMap {
    /// Generates a usable source map, capable of both tracing and stringifying.
    #[turbo_tasks::function]
    fn generate_source_map(self: Vc<Self>) -> Vc<FileContent>;

    /// Returns an individual section of the larger source map, if found.
    #[turbo_tasks::function]
    fn by_section(self: Vc<Self>, _section: RcStr) -> Vc<FileContent> {
        FileContent::NotFound.cell()
    }
}

/// Implements the source map specification as either a decoded or sectioned sourcemap.
///
/// - A "decoded" map represents a source map as if it was came out of a JSON
/// decode.
/// - A "sectioned" source map is a tree of many [SourceMap]
/// covering regions of an output file.
///
/// The distinction between the source map spec's [sourcemap::Index] and our
/// [SourceMap::Sectioned] is whether the sections are represented with Vcs
/// pointers.
#[turbo_tasks::value(shared, cell = "new", eq = "manual")]
#[derive(Debug)]
pub struct SourceMap {
    /// A decoded source map contains no Vcs.
    #[turbo_tasks(trace_ignore)]
    map: Arc<CrateMapWrapper>,
}
impl Eq for SourceMap {}
impl PartialEq for SourceMap {
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.map, &other.map)
    }
}

#[turbo_tasks::value(transparent)]
pub struct OptionSourceMap(Option<SourceMap>);

#[turbo_tasks::value_impl]
impl OptionSourceMap {
    #[turbo_tasks::function]
    pub fn none() -> Vc<Self> {
        Vc::cell(None)
    }
}

impl OptionSourceMap {
    pub fn none_resolved() -> ResolvedVc<Self> {
        ResolvedVc::cell(None)
    }
}

/// A token represents a mapping in a source map.
///
/// It may either be Synthetic, meaning it was generated by some build tool and doesn't
/// represent a location in a user-authored source file, or it is Original, meaning it represents a
/// real location in source file.
#[turbo_tasks::value]
#[derive(Clone, Debug)]
pub enum Token {
    Synthetic(SyntheticToken),
    Original(OriginalToken),
}

#[turbo_tasks::value]
#[derive(Clone, Debug)]
pub struct TokenWithSource {
    pub token: Token,
    pub source_content: Option<ResolvedVc<Box<dyn Source>>>,
}

/// A SyntheticToken represents a region of the generated file that was created
/// by some build tool.
#[turbo_tasks::value]
#[derive(Clone, Debug)]
pub struct SyntheticToken {
    pub generated_line: u32,
    pub generated_column: u32,
    pub guessed_original_file: Option<RcStr>,
}

/// An OriginalToken represents a region of the generated file that exists in
/// user-authored source file.
#[turbo_tasks::value]
#[derive(Clone, Debug)]
pub struct OriginalToken {
    pub generated_line: u32,
    pub generated_column: u32,
    pub original_file: RcStr,
    pub original_line: u32,
    pub original_column: u32,
    pub name: Option<RcStr>,
}

impl Token {
    pub fn generated_line(&self) -> u32 {
        match self {
            Self::Original(t) => t.generated_line,
            Self::Synthetic(t) => t.generated_line,
        }
    }

    pub fn generated_column(&self) -> u32 {
        match self {
            Self::Original(t) => t.generated_column,
            Self::Synthetic(t) => t.generated_column,
        }
    }

    pub fn with_offset(&self, line_offset: u32, column_offset: u32) -> Self {
        match self {
            Self::Original(t) => Self::Original(OriginalToken {
                generated_line: t.generated_line + line_offset,
                generated_column: if t.generated_line == 0 {
                    t.generated_column + column_offset
                } else {
                    t.generated_column
                },
                original_file: t.original_file.clone(),
                original_line: t.original_line,
                original_column: t.original_column,
                name: t.name.clone(),
            }),
            Self::Synthetic(t) => Self::Synthetic(SyntheticToken {
                generated_line: t.generated_line + line_offset,
                generated_column: if t.generated_line == 0 {
                    t.generated_column + column_offset
                } else {
                    t.generated_column
                },
                guessed_original_file: t.guessed_original_file.clone(),
            }),
        }
    }
}

impl From<swc_sourcemap::Token<'_>> for Token {
    fn from(t: swc_sourcemap::Token) -> Self {
        if t.has_source() {
            Token::Original(OriginalToken {
                generated_line: t.get_dst_line(),
                generated_column: t.get_dst_col(),
                original_file: RcStr::from(
                    t.get_source()
                        .expect("already checked token has source")
                        .clone(),
                ),
                original_line: t.get_src_line(),
                original_column: t.get_src_col(),
                name: t.get_name().cloned().map(RcStr::from),
            })
        } else {
            Token::Synthetic(SyntheticToken {
                generated_line: t.get_dst_line(),
                generated_column: t.get_dst_col(),
                guessed_original_file: None,
            })
        }
    }
}

impl TryInto<swc_sourcemap::RawToken> for Token {
    type Error = std::num::ParseIntError;

    fn try_into(self) -> Result<swc_sourcemap::RawToken, Self::Error> {
        Ok(match self {
            Self::Original(t) => swc_sourcemap::RawToken {
                dst_col: t.generated_column,
                dst_line: t.generated_line,
                name_id: match t.name {
                    None => SOURCEMAP_CRATE_NONE_U32,
                    Some(name) => name.parse()?,
                },
                src_col: t.original_column,
                src_line: t.original_line,
                src_id: t.original_file.parse()?,
                is_range: false,
            },
            Self::Synthetic(t) => swc_sourcemap::RawToken {
                dst_col: t.generated_column,
                dst_line: t.generated_line,
                name_id: SOURCEMAP_CRATE_NONE_U32,
                src_col: SOURCEMAP_CRATE_NONE_U32,
                src_line: SOURCEMAP_CRATE_NONE_U32,
                src_id: SOURCEMAP_CRATE_NONE_U32,
                is_range: false,
            },
        })
    }
}

impl SourceMap {
    /// Creates a new SourceMap::Decoded Vc out of a [RegularMap] instance.
    fn new_regular(map: RegularMap) -> Self {
        Self::new_decoded(DecodedMap::Regular(map))
    }

    /// Creates a new SourceMap::Decoded Vc out of a [DecodedMap] instance.
    fn new_decoded(map: DecodedMap) -> Self {
        SourceMap {
            map: Arc::new(CrateMapWrapper(map)),
        }
    }

    pub fn new_from_rope(content: &Rope) -> Result<Option<Self>> {
        let Ok(map) = DecodedMap::from_reader(content.read()) else {
            return Ok(None);
        };
        Ok(Some(SourceMap::new_decoded(map)))
    }
}

#[turbo_tasks::value_impl]
impl SourceMap {
    /// This function should be used sparingly to reduce memory usage, only in cold code paths
    /// (issue resolving, etc).
    #[turbo_tasks::function]
    pub async fn new_from_rope_cached(content: Vc<FileContent>) -> Result<Vc<OptionSourceMap>> {
        let content = content.await?;
        let Some(content) = content.as_content() else {
            return Ok(OptionSourceMap::none());
        };
        Ok(Vc::cell(SourceMap::new_from_rope(content.content())?))
    }
}

impl SourceMap {
    pub fn to_source_map(&self) -> Arc<CrateMapWrapper> {
        self.map.clone()
    }
}

static EMPTY_SOURCE_MAP_ROPE: Lazy<Rope> =
    Lazy::new(|| Rope::from(r#"{"version":3,"sources":[],"names":[],"mappings":"A"}"#));

impl SourceMap {
    /// A source map that contains no actual source location information (no
    /// `sources`, no mappings that point into a source). This is used to tell
    /// Chrome that the generated code starting at a particular offset is no
    /// longer part of the previous section's mappings.
    pub fn empty() -> Self {
        let mut builder = SourceMapBuilder::new(None);
        builder.add(0, 0, 0, 0, None, None, false);
        SourceMap::new_regular(builder.into_sourcemap())
    }

    /// A source map that contains no actual source location information (no
    /// `sources`, no mappings that point into a source). This is used to tell
    /// Chrome that the generated code starting at a particular offset is no
    /// longer part of the previous section's mappings.
    pub fn empty_rope() -> Rope {
        EMPTY_SOURCE_MAP_ROPE.clone()
    }

    pub fn sections_to_rope(
        sections: impl IntoIterator<Item = (SourcePos, Rope)>,
        debug_id: Option<RcStr>,
    ) -> Rope {
        let mut sections = sections.into_iter().peekable();

        let mut first = sections.next();
        if let Some((offset, map)) = &mut first
            && sections.peek().is_none()
            && *offset == (0, 0)
            && debug_id.is_none()
        {
            // There is just a single sourcemap that starts at the beginning of the file.
            return std::mem::take(map);
        }

        // My kingdom for a decent dedent macro with interpolation!
        // NOTE: The empty `sources` array is technically incorrect, but there is a bug
        // in Node.js that requires sectioned source maps to have a `sources` array.
        let mut rope = RopeBuilder::from(
            r#"{
  "version": 3,
  "sources": [],
"#,
        );
        if let Some(debug_id) = debug_id {
            writeln!(rope, r#"  "debugId": "{debug_id}","#).unwrap();
        }
        rope += "  \"sections\": [";

        let mut first_section = true;
        for (offset, section_map) in first.into_iter().chain(sections) {
            if !first_section {
                rope += ",";
            }
            first_section = false;

            write!(
                rope,
                r#"
    {{"offset": {{"line": {}, "column": {}}}, "map": "#,
                offset.line, offset.column,
            )
            .unwrap();

            rope += &section_map;

            rope += "}";
        }

        rope += "]";

        rope += "\n}";

        rope.build()
    }

    /// Stringifies the source map into JSON bytes.
    pub fn to_rope(&self) -> Result<Rope> {
        let mut bytes = vec![];
        self.map.0.to_writer(&mut bytes)?;
        Ok(Rope::from(bytes))
    }

    /// Traces a generated line/column into an mapping token representing either
    /// synthetic code or user-authored original code.
    pub fn lookup_token(&self, line: u32, column: u32) -> Token {
        let (token, _) = self.lookup_token_and_source_internal(line, column, true);
        token
    }

    /// Traces a generated line/column into an mapping token representing either
    /// synthetic code or user-authored original code.
    pub async fn lookup_token_and_source(&self, line: u32, column: u32) -> Result<TokenWithSource> {
        let (token, content) = self.lookup_token_and_source_internal(line, column, true);
        Ok(TokenWithSource {
            token,
            source_content: match content {
                Some(v) => Some(v.to_resolved().await?),
                None => None,
            },
        })
    }

    pub async fn with_resolved_sources(&self, origin: FileSystemPath) -> Result<Self> {
        async fn resolve_source(
            source_request: BytesStr,
            source_content: Option<BytesStr>,
            origin: FileSystemPath,
        ) -> Result<(BytesStr, BytesStr)> {
            Ok(
                if let Some(path) = origin.parent().try_join(&source_request) {
                    let path_str = path.value_to_string().await?;
                    let source = format!("{SOURCE_URL_PROTOCOL}///{path_str}");
                    let source_content = if let Some(source_content) = source_content {
                        source_content
                    } else if let FileContent::Content(file) = &*path.read().await? {
                        let text = file.content().to_str()?;
                        text.to_string().into()
                    } else {
                        format!("unable to read source {path_str}").into()
                    };
                    (source.into(), source_content)
                } else {
                    let origin_str = origin.value_to_string().await?;
                    static INVALID_REGEX: Lazy<Regex> =
                        Lazy::new(|| Regex::new(r#"(?:^|/)(?:\.\.?(?:/|$))+"#).unwrap());
                    let source = INVALID_REGEX
                        .replace_all(&source_request, |s: &regex::Captures<'_>| {
                            s[0].replace('.', "_")
                        });
                    let source = format!("{SOURCE_URL_PROTOCOL}///{origin_str}/{source}");
                    let source_content = source_content.unwrap_or_else(|| {
                        format!(
                            "unable to access {source_request} in {origin_str} (it's leaving the \
                             filesystem root)"
                        )
                        .into()
                    });
                    (source.into(), source_content)
                },
            )
        }
        async fn regular_map_with_resolved_sources(
            map: &RegularMapWrapper,
            origin: FileSystemPath,
        ) -> Result<RegularMap> {
            let map = &map.0;
            let file = map.get_file().cloned();
            let tokens = map.tokens().map(|t| t.get_raw_token()).collect();
            let names = map.names().cloned().collect();
            let count = map.get_source_count() as usize;
            let sources = map.sources().cloned().collect::<Vec<_>>();
            let source_contents = map
                .source_contents()
                .map(|s| s.cloned())
                .collect::<Vec<_>>();
            let mut new_sources = Vec::with_capacity(count);
            let mut new_source_contents = Vec::with_capacity(count);
            for (source, source_content) in sources.into_iter().zip(source_contents.into_iter()) {
                let (source, source_content) =
                    resolve_source(source, source_content, origin.clone()).await?;
                new_sources.push(source);
                new_source_contents.push(Some(source_content));
            }
            let mut map =
                RegularMap::new(file, tokens, names, new_sources, Some(new_source_contents));

            add_default_ignore_list(&mut map);

            Ok(map)
        }
        async fn decoded_map_with_resolved_sources(
            map: &CrateMapWrapper,
            origin: FileSystemPath,
        ) -> Result<CrateMapWrapper> {
            Ok(CrateMapWrapper(match &map.0 {
                DecodedMap::Regular(map) => {
                    let map = RegularMapWrapper::ref_cast(map);
                    DecodedMap::Regular(regular_map_with_resolved_sources(map, origin).await?)
                }
                DecodedMap::Index(map) => {
                    let count = map.get_section_count() as usize;
                    let file = map.get_file().cloned();
                    let sections = map
                        .sections()
                        .filter_map(|section| {
                            section
                                .get_sourcemap()
                                .map(|s| (section.get_offset(), CrateMapWrapper::ref_cast(s)))
                        })
                        .collect::<Vec<_>>();
                    let sections = sections
                        .into_iter()
                        .map(|(offset, map)| {
                            let origin = origin.clone();
                            async move {
                                Ok((
                                    offset,
                                    Box::pin(decoded_map_with_resolved_sources(
                                        map,
                                        origin.clone(),
                                    ))
                                    .await?,
                                ))
                            }
                        })
                        .try_join()
                        .await?;
                    let mut new_sections = Vec::with_capacity(count);
                    for (offset, map) in sections {
                        new_sections.push(swc_sourcemap::SourceMapSection::new(
                            offset,
                            // Urls are deprecated and we don't accept them
                            None,
                            Some(map.0),
                        ));
                    }
                    DecodedMap::Index(SourceMapIndex::new(file, new_sections))
                }
                DecodedMap::Hermes(_) => {
                    todo!("hermes source maps are not implemented");
                }
            }))
        }

        let map = Box::pin(decoded_map_with_resolved_sources(&self.map, origin)).await?;
        Ok(Self::new_decoded(map.0))
    }
}

#[turbo_tasks::function]
fn sourcemap_content_fs_root() -> Vc<FileSystemPath> {
    VirtualFileSystem::new_with_name(rcstr!("sourcemap-content")).root()
}

#[turbo_tasks::function]
async fn sourcemap_content_source(path: RcStr, content: RcStr) -> Result<Vc<Box<dyn Source>>> {
    let path = sourcemap_content_fs_root().await?.join(&path)?;
    let content = AssetContent::file(FileContent::new(File::from(content)).cell());
    Ok(Vc::upcast(VirtualSource::new(path, content)))
}

impl SourceMap {
    fn lookup_token_and_source_internal(
        &self,
        line: u32,
        column: u32,
        need_source_content: bool,
    ) -> (Token, Option<Vc<Box<dyn Source>>>) {
        let mut content: Option<Vc<Box<dyn Source>>> = None;

        let token: Token = {
            let map = &self.map;

            let tok = map.lookup_token(line, column);
            let mut token = tok.map(Token::from).unwrap_or_else(|| {
                Token::Synthetic(SyntheticToken {
                    generated_line: line,
                    generated_column: column,
                    guessed_original_file: None,
                })
            });

            if let Token::Synthetic(SyntheticToken {
                guessed_original_file,
                ..
            }) = &mut token
                && let DecodedMap::Regular(map) = &map.0
                && map.get_source_count() == 1
            {
                let source = map.sources().next().unwrap().clone();
                *guessed_original_file = Some(RcStr::from(source));
            }

            if need_source_content
                && content.is_none()
                && let Some(map) = map.as_regular_source_map()
            {
                content = tok.and_then(|tok| {
                    let src_id = tok.get_src_id();

                    let name = map.get_source(src_id);
                    let content = map.get_source_contents(src_id);

                    let (name, content) = name.zip(content)?;
                    Some(sourcemap_content_source(
                        name.clone().into(),
                        content.clone().into(),
                    ))
                });
            }

            token
        };

        (token, content)
    }
}

impl SourceMap {
    pub fn tokens(&self) -> impl Iterator<Item = Token> + '_ {
        let map = &self.map;

        fn regular_map_to_tokens(
            map: &RegularMap,
            offset_line: u32,
            offset_column: u32,
        ) -> impl Iterator<Item = Token> + '_ {
            map.tokens()
                .map(move |t| Token::from(t).with_offset(offset_line, offset_column))
        }

        fn index_map_to_tokens(
            map: &SourceMapIndex,
            offset_line: u32,
            offset_column: u32,
        ) -> impl Iterator<Item = Token> + '_ {
            map.sections().flat_map(move |section| {
                let (line, col) = section.get_offset();
                let offset_line = offset_line + line;
                let offset_column = if line == 0 { offset_column + col } else { col };
                if let Some(source_map) = section.get_sourcemap() {
                    Either::Left(Box::new(decoded_map_to_tokens(
                        source_map,
                        offset_line,
                        offset_column,
                    )) as Box<dyn Iterator<Item = Token>>)
                } else {
                    Either::Right(std::iter::empty())
                }
            })
        }

        fn decoded_map_to_tokens(
            map: &DecodedMap,
            offset_line: u32,
            offset_column: u32,
        ) -> impl Iterator<Item = Token> + '_ {
            match map {
                DecodedMap::Regular(map) => {
                    Either::Left(regular_map_to_tokens(map, offset_line, offset_column))
                }
                DecodedMap::Index(map) => {
                    Either::Right(index_map_to_tokens(map, offset_line, offset_column))
                }
                DecodedMap::Hermes(_) => {
                    todo!("hermes source maps are not implemented");
                }
            }
        }

        decoded_map_to_tokens(&map.0, 0, 0)
    }
}

#[turbo_tasks::value_impl]
impl GenerateSourceMap for SourceMap {
    #[turbo_tasks::function]
    fn generate_source_map(&self) -> Result<Vc<FileContent>> {
        Ok(FileContent::Content(File::from(self.to_rope()?)).cell())
    }
}

/// Wraps the DecodedMap struct so that it is Sync and Send.
///
/// # Safety
///
/// Must not use per line access to the SourceMap, as it is not thread safe.
#[derive(Debug, RefCast)]
#[repr(transparent)]
pub struct CrateMapWrapper(DecodedMap);

// Safety: DecodedMap contains a raw pointer, which isn't Send, which is
// required to cache in a Vc. So, we have wrap it in 4 layers of cruft to do it.
unsafe impl Send for CrateMapWrapper {}
unsafe impl Sync for CrateMapWrapper {}

/// Wraps the RegularMap struct so that it is Sync and Send.
///
/// # Safety
///
/// Must not use per line access to the SourceMap, as it is not thread safe.
#[derive(Debug, RefCast)]
#[repr(transparent)]
pub struct RegularMapWrapper(RegularMap);

// Safety: RegularMap contains a raw pointer, which isn't Send, which is
// required to cache in a Vc. So, we have wrap it in 4 layers of cruft to do it.
unsafe impl Send for RegularMapWrapper {}
unsafe impl Sync for RegularMapWrapper {}

impl CrateMapWrapper {
    pub fn as_regular_source_map(&self) -> Option<Cow<'_, RegularMap>> {
        match &self.0 {
            DecodedMap::Regular(m) => Some(Cow::Borrowed(m)),
            DecodedMap::Index(m) => m.flatten().map(Cow::Owned).ok(),
            _ => None,
        }
    }
}

impl Deref for CrateMapWrapper {
    type Target = DecodedMap;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Serialize for CrateMapWrapper {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        use serde::ser::Error;
        let mut bytes = vec![];
        self.0.to_writer(&mut bytes).map_err(Error::custom)?;
        serializer.serialize_bytes(bytes.as_slice())
    }
}

impl<'de> Deserialize<'de> for CrateMapWrapper {
    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
        use serde::de::Error;
        let bytes = <&[u8]>::deserialize(deserializer)?;
        let map = DecodedMap::from_reader(bytes).map_err(Error::custom)?;
        Ok(CrateMapWrapper(map))
    }
}

impl Encode for CrateMapWrapper {
    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
        let mut bytes = Vec::new();
        self.0
            .to_writer(&mut bytes)
            .map_err(|e| EncodeError::OtherString(e.to_string()))?;
        bytes.encode(encoder)
    }
}

impl<Context> Decode<Context> for CrateMapWrapper {
    fn decode<D: Decoder<Context = Context>>(decoder: &mut D) -> Result<Self, DecodeError> {
        let bytes = Vec::<u8>::decode(decoder)?;
        let map = DecodedMap::from_reader(&*bytes)
            .map_err(|e| DecodeError::OtherString(e.to_string()))?;
        Ok(CrateMapWrapper(map))
    }
}

bincode::impl_borrow_decode!(CrateMapWrapper);
